import numpy as np
import random

# 四旋翼参数
MASS = 1.0         # kg
GRAVITY = 9.81     # m/s^2
DT = 0.01          # 时间步长

# 飞行参数
WAYPOINT_RADIUS = 0.1  # 到达航路点的判定半径

# 航路点参数
MAX_WAYPOINTS = 3     # 每架无人机的最大航路点数
MAX_DISTANCE = 100    # 最大航路点距离

# 速度控制参数
MAX_SPEED = 15.0      # 最大飞行速度 (m/s)
APPROACH_DISTANCE = 2.0  # 接近航路点时开始减速的距离
INERTIA = np.array([0.1, 0.1, 0.1])  # 转动惯量
# PID 控制器
class PID:
    def __init__(self, kp, ki, kd, setpoint=0):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.setpoint = setpoint
        self.integral = 0
        self.last_error = 0

    def update(self, value, dt):
        error = self.setpoint - value
        self.integral += error * dt
        derivative = (error - self.last_error) / dt
        self.last_error = error
        return self.kp * error + self.ki * self.integral + self.kd * derivative

    def set_setpoint(self, setpoint):
        self.setpoint = setpoint

# 无人机类
class Drone:
    def __init__(self, drone_id, initial_pos):
        self.id = drone_id
        
        # 状态变量
        self.pos = initial_pos.copy()  # 位置 x, y, z
        self.vel = np.array([0.0, 0.0, 0.0])   # 速度 vx, vy, vz
        self.att = np.array([0.0, 0.0, 0.0])   # 姿态角 roll, pitch, yaw
        self.ang_vel = np.array([0.0, 0.0, 0.0])  # 角速度
        
        # 随机生成航路点
        num_waypoints = random.randint(1, MAX_WAYPOINTS)
        self.waypoints = []
        for _ in range(num_waypoints):
            x = random.uniform(-MAX_DISTANCE, MAX_DISTANCE)
            y = random.uniform(-MAX_DISTANCE, MAX_DISTANCE)
            z = random.uniform(10, 50)  # 随机高度
            self.waypoints.append(np.array([x, y, z]))
        
        self.current_waypoint_index = 0
        self.active = True
        
        # PID控制器
        self.pid_vx = PID(2.0, 0.0, 0.5)
        self.pid_vy = PID(2.0, 0.0, 0.5)
        self.pid_vz = PID(2.0, 0.0, 0.5)
        self.pid_roll = PID(2.0, 0.0, 0.5)
        self.pid_pitch = PID(2.0, 0.0, 0.5)
        self.pid_yaw = PID(2.0, 0.0, 0.5)
        
        # 运动计算相关
        self.acc = np.array([0.0, 0.0, 0.0])
        self.calculated_att = np.array([0.0, 0.0, 0.0])
        self.filtered_pitch = 0.0
        self.filtered_yaw = 0.0
    
    def calculate_velocity_setpoint(self, target_pos):
        """计算速度设定点，考虑最大速度和接近目标时的减速"""
        direction = target_pos - self.pos
        distance = np.linalg.norm(direction)
        
        if distance < 0.01:  # 防止除以零
            return np.zeros(3)
        
        direction_unit = direction / distance
        
        if distance > APPROACH_DISTANCE:
            return direction_unit * MAX_SPEED
        else:
            speed_factor = distance / APPROACH_DISTANCE
            return direction_unit * MAX_SPEED * speed_factor
    
    def calculate_attitude_from_motion(self):
        """根据速度和加速度计算平滑的姿态信息"""
        MAX_PITCH_ANGLE = 20  # 最大俯仰角(度)
        MAX_ROLL_ANGLE = 30   # 最大滚转角(度)
        alpha = 0.2  # 低通滤波系数
        
        # 计算航向角 (yaw)
        if np.linalg.norm(self.vel[:2]) > 0.1:
            new_yaw = np.arctan2(self.vel[1], self.vel[0]) * 57.3  # 转换为度
            if abs(new_yaw - self.filtered_yaw) > 180:
                if new_yaw > self.filtered_yaw:
                    new_yaw -= 360
                else:
                    new_yaw += 360
            self.filtered_yaw = alpha * new_yaw + (1 - alpha) * self.filtered_yaw
        else:
            self.filtered_yaw = self.calculated_att[2]
        
        # 计算俯仰角 (pitch)
        yaw_rad = np.radians(self.filtered_yaw)
        forward_accel = self.acc[0] * np.cos(yaw_rad) + self.acc[1] * np.sin(yaw_rad)
        forward_speed = self.vel[0] * np.cos(yaw_rad) + self.vel[1] * np.sin(yaw_rad)
        pitch_direction = np.sign(forward_accel) if abs(forward_accel) > 0.1 else 0
        speed_factor = np.clip(abs(forward_speed) / MAX_SPEED, 0.0, 1.0)
        target_pitch = pitch_direction * speed_factor * MAX_PITCH_ANGLE
        self.filtered_pitch = alpha * target_pitch + (1 - alpha) * self.filtered_pitch
        
        # 计算滚转角 (roll)
        if np.linalg.norm(self.vel[:2]) > 0.5:
            lateral_accel = -self.acc[0] * np.sin(yaw_rad) + self.acc[1] * np.cos(yaw_rad)
            target_roll = np.clip(np.arctan(lateral_accel / 9.81) * 57.3, -MAX_ROLL_ANGLE, MAX_ROLL_ANGLE)
        else:
            target_roll = 0.0
        
        self.calculated_att = np.array([
            alpha * target_roll + (1 - alpha) * self.calculated_att[0],
            self.filtered_pitch,
            self.filtered_yaw
        ])
        
        return self.calculated_att
    
    def update_waypoint(self):
        current_waypoint = self.waypoints[self.current_waypoint_index]
        distance = np.linalg.norm(self.pos - current_waypoint)
        
        if distance < WAYPOINT_RADIUS:
            self.current_waypoint_index += 1
            if self.current_waypoint_index >= len(self.waypoints):
                self.active = False
    
    def update(self):
        if not self.active:
            return False
        
        # 更新航路点
        self.update_waypoint()
        if not self.active:
            return False
        
        # 计算速度设定点
        target_pos = self.waypoints[self.current_waypoint_index]
        velocity_setpoint = self.calculate_velocity_setpoint(target_pos)
        self.pid_vx.set_setpoint(velocity_setpoint[0])
        self.pid_vy.set_setpoint(velocity_setpoint[1])
        self.pid_vz.set_setpoint(velocity_setpoint[2])
        
        # 速度PID控制器输出
        fx = self.pid_vx.update(self.vel[0], DT)
        fy = self.pid_vy.update(self.vel[1], DT)
        fz = self.pid_vz.update(self.vel[2], DT) + MASS * GRAVITY

        # 计算姿态控制力矩
        tau_roll = self.pid_roll.update(self.att[0], DT)
        tau_pitch = self.pid_pitch.update(self.att[1], DT)
        tau_yaw = self.pid_yaw.update(self.att[2], DT)

        # 合力和力矩
        force = np.array([fx, fy, fz])
        torque = np.array([tau_roll, tau_pitch, tau_yaw])

        # 加速度
        self.acc = force / MASS - np.array([0, 0, GRAVITY])
        ang_acc = torque / INERTIA

        # 更新速度和位置
        self.vel += self.acc * DT
        self.pos += self.vel * DT
        
        # 防止降到地面以下
        if self.pos[2] < 0:
            self.pos[2] = 0
            self.vel[2] = 0

        # 更新角速度和姿态
        self.ang_vel += ang_acc * DT
        self.att += self.ang_vel * DT

        # 计算基于运动的姿态信息
        self.calculate_attitude_from_motion()
        
        return True

def main():
    # 创建无人机
    initial_pos = np.array([0.0, 0.0, 0.0])
    drone = Drone(0, initial_pos)
    
    # 模拟循环
    sim_time = 0
    while drone.active:
        drone.update()
        sim_time += DT
        
        # 可以在这里打印状态或记录数据
        if int(sim_time * 10) % 10 == 0:  # 每0.1秒打印一次
            print(f"Time: {sim_time:.2f}s, Position: {drone.pos}, Waypoint: {drone.current_waypoint_index+1}/{len(drone.waypoints)}")
    
    print("Drone has completed all waypoints!")

if __name__ == "__main__":
    main()